#!/usr/bin/env python3

'''	Loads the dataset 2a of the BCI Competition IV
available on http://bnci-horizon-2020.eu/database/data-sets
'''
import os
import numpy as np
import scipy.io as sio
from sklearn.model_selection import train_test_split
from keras.utils.np_utils import to_categorical
__author__ = "Michael Hersche and Tino Rellstab"
__email__ = "herschmi@ethz.ch,tinor@ethz.ch"
data_path = "/media/brainseek/dataset/9subjects/rawdata/"
save_path = "/media/brainseek/dataset/9subjects/LB_2s_250Hz/"

def get_data(subject, training):
    '''	Loads the dataset 2a of the BCI Competition IV
    available on http://bnci-horizon-2020.eu/database/data-sets
    Keyword arguments:
    subject -- number of subject in [1, .. ,9]
    training -- if True, load training data
                if False, load testing data

    Return:	data_return 	numpy matrix 	size = NO_valid_trial x 22 x 1750
            class_return 	numpy matrix 	size = NO_valid_trial
    '''
    NO_channels = 22
    NO_tests = 6 * 48
    Window_Length = 7 * 250

    class_return = np.zeros(NO_tests)
    data_return = np.zeros((NO_tests, NO_channels, Window_Length))

    NO_valid_trial = 0
    if training:
        a = sio.loadmat(data_path + 'A0' + str(subject) + 'T.mat')
    else:
        a = sio.loadmat(data_path + 'A0' + str(subject) + 'E.mat')
    a_data = a['data']
    for ii in range(0, a_data.size):
        a_data1 = a_data[0, ii]
        a_data2 = [a_data1[0, 0]]
        a_data3 = a_data2[0]
        a_X = a_data3[0]
        a_trial = a_data3[1]
        a_y = a_data3[2]
        a_fs = a_data3[3]
        a_classes = a_data3[4]
        a_artifacts = a_data3[5]
        a_gender = a_data3[6]
        a_age = a_data3[7]
        for trial in range(0, a_trial.size):
            if (a_artifacts[trial] == 0):
                data_return[NO_valid_trial, :, :] = np.transpose(
                    a_X[int(a_trial[trial]):(int(a_trial[trial]) + Window_Length), :22])
                class_return[NO_valid_trial] = int(a_y[trial])
                NO_valid_trial += 1
    print('data loaded' + str(subject))
    return data_return[0:NO_valid_trial, :, :], class_return[0:NO_valid_trial]

if __name__ == '__main__':
    data_total_raw = np.empty((0,12, 1750))
    subjects_good = np.arange(1, 10)

    if not os.path.exists(save_path):
        os.mkdir(save_path)
    for subs in subjects_good:
        res = get_data(subs,training=True)
        for num in np.arange(0, res[0].shape[0]):
            if res[1][num] == 1:  ##左手运动想象分类标签为1
                data_left = res[0][num][[1,5,2,4,7,11,8,10,13,17,14,16]]
                data_total_raw = np.concatenate((data_total_raw, data_left.reshape(1,12,1750)), axis=0)

    data_total_raw  = data_total_raw.transpose(0, 2, 1)
    left_data_total = data_total_raw[:,750:1250,:]
    base_data_total = data_total_raw[:,0:500,:]
    # get the data and labels 0:base  1:left
    left_label_total = np.ones([1,left_data_total.shape[0]])
    base_label_total = np.zeros([1,base_data_total.shape[0]])

    data_total  = np.concatenate((left_data_total, base_data_total), axis=0)
    label_total = np.concatenate((left_label_total, base_label_total), axis=1)
    labels = to_categorical(label_total[0])  # one-hot
    train_data_ori, test_data_ori, train_label_ori, test_label_ori =    train_test_split(
                                                                        data_total,
                                                                        labels,
                                                                        test_size=0.2,
                                                                        random_state=42)

    train_data_total = np.empty((0, 500, 2))
    test_data_total  = np.empty((0, 500, 2))
    train_label_total = np.empty((0, 2))
    test_label_total  = np.empty((0, 2))
    for index in range(0, 12, 2):
        train_data_total = np.concatenate((train_data_total, train_data_ori[:,:,index:index+2]), axis=0)
        test_data_total  = np.concatenate((test_data_total, test_data_ori[:,:,index:index+2]), axis=0)
        train_label_total  = np.concatenate((train_label_total,train_label_ori ), axis=0)
        test_label_total   = np.concatenate((test_label_total, test_label_ori ), axis=0)
    np.save(os.path.join(save_path, "train_data_total" ),  train_data_total, allow_pickle=True)
    np.save(os.path.join(save_path, "test_data_total" ),   test_data_total, allow_pickle=True)
    np.save(os.path.join(save_path, "train_label_total" ), train_label_total, allow_pickle=True)
    np.save(os.path.join(save_path, "test_label_total" ),  test_label_total, allow_pickle=True)
    print(train_data_total.shape)
    print(test_data_total.shape)
    print(train_label_total.shape)
    print(test_label_total.shape)